package net.aiion.weave.impl.painters.transformation;

import java.util.ArrayList;
import java.util.Collection;

import net.aiion.weave.asm.Label;
import net.aiion.weave.asm.Opcodes;
import net.aiion.weave.asm.Type;
import net.aiion.weave.spi.bytecode.Bytecode;
import net.aiion.weave.spi.bytecode.BytecodeContext;
import net.aiion.weave.spi.bytecode.BytecodeException;
import net.aiion.weave.spi.bytecode.BytecodeGenerator;
import net.aiion.weave.spi.bytecode.events.AnnotationEndBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.AnnotationStartBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.AnnotationValueBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.BytecodeEvent;
import net.aiion.weave.spi.bytecode.events.ClassEndBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.ClassStartBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.MethodEndBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.MethodStartBytecodeEvent;
import net.aiion.weave.spi.bytecode.util.BytecodeTransformation;
import net.aiion.weave.spi.painters.PaintBrush;
import net.aiion.weave.spi.painters.PaintContext;
import net.aiion.weave.spi.painters.Painter;
import net.aiion.weave.spi.painters.annotations.Paints;

//TODO Add inheritance support (also respect annotations of implemented interfaces).
//TODO Add real support for multiple @Paints annotated methods. Currently, a call to PainterContext.forwardPainting() skips all other @Injects annotated methods for the same class.
public class PainterTransformation extends BytecodeTransformation {
    private static final String PAINT_CONTEXT_INTERNAL_NAME = Type.getInternalName(PaintContext.class);
    private static final String PAINT_BRUSH_INTERNAL_NAME = Type.getInternalName(PaintBrush.class);

    private static final String PAINTER_PAINT_COMPONENT_DESC = "(L" + PAINT_CONTEXT_INTERNAL_NAME + ";)L"
        + PAINT_BRUSH_INTERNAL_NAME + ";";

    private String className;

    private String methodName;
    private String methodDescriptor;

    private Type paintedType;

    private boolean inPaintsAnnotation;

    private Collection<BytecodeEvent> paintComponentBytecodeEvents;
    private BytecodeGenerator paintComponentBytecodeGenerator;

    public PainterTransformation() {
    }

    @Override
    public boolean testRetain(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        return aBytecodeEvent.getType() == BytecodeEvent.CLASS_START
            && !aBytecodeEvent.asClassStart().containsInterface(Painter.class);
    }

    @Override
    public boolean testProcess(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        return aBytecodeEvent.getType() == BytecodeEvent.ANNOTATION_START
            && aBytecodeEvent.asAnnotationStart().equalsType(Paints.class);
    }

    @Override
    public boolean testSkip(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        return aBytecodeEvent.getType() == BytecodeEvent.CLASS_START
            && (aBytecodeEvent.asClassStart().isInterface() || aBytecodeEvent.asClassStart().containsInterface(
                Painter.class));
    }

    @Override
    public void handleBytecodeEvent(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        switch (aBytecodeEvent.getType()) {
        case BytecodeEvent.CLASS_START:
            handleClassStartBytecode(aBytecodeContext, aBytecodeEvent.asClassStart());
            break;
        case BytecodeEvent.METHOD_START:
            handleMethodStartBytecode(aBytecodeContext, aBytecodeEvent.asMethodStart());
            break;
        case BytecodeEvent.ANNOTATION_START:
            handleAnnotationStartBytecode(aBytecodeContext, aBytecodeEvent.asAnnotationStart());
            break;
        case BytecodeEvent.ANNOTATION_VALUE:
            handleAnnotationValueBytecode(aBytecodeContext, aBytecodeEvent.asAnnotationValue());
            break;
        case BytecodeEvent.ANNOTATION_END:
            handleAnnotationEndBytecode(aBytecodeContext, aBytecodeEvent.asAnnotationEnd());
            break;
        case BytecodeEvent.METHOD_END:
            handleMethodEndBytecode(aBytecodeContext, aBytecodeEvent.asMethodEnd());
            break;
        case BytecodeEvent.CLASS_END:
            handleClassEndBytecode(aBytecodeContext, aBytecodeEvent.asClassEnd());
            break;
        default:
            break;
        }

        aBytecodeContext.forwardBytecodeEvent(aBytecodeEvent);
    }

    private void handleClassStartBytecode(final BytecodeContext aBytecodeContext,
        final ClassStartBytecodeEvent aClassStartBytecodeEvent) {
        className = aClassStartBytecodeEvent.getName();
        aClassStartBytecodeEvent.addInterface(Painter.class);

        preparePaintComponent();
    }

    private void handleMethodStartBytecode(final BytecodeContext aBytecodeContext,
        final MethodStartBytecodeEvent aMethodStartBytecodeEvent) {
        methodName = aMethodStartBytecodeEvent.getName();
        methodDescriptor = aMethodStartBytecodeEvent.getDescriptor();
    }

    private void handleAnnotationStartBytecode(final BytecodeContext aBytecodeContext,
        final AnnotationStartBytecodeEvent aAnnotationStartBytecodeEvent) {
        if (methodName != null && aAnnotationStartBytecodeEvent.equalsType(Paints.class)) {
            inPaintsAnnotation = true;

            testConverterMethod();
        }
    }

    public void testConverterMethod() {
        if (!methodDescriptor.equals(PAINTER_PAINT_COMPONENT_DESC)) {
            throw new BytecodeException("A @Paints annotated method must accept exactly one argument of type "
                + PaintContext.class.getName() + " and must return " + PaintBrush.class.getName() + ".");
        }
    }

    private void handleAnnotationValueBytecode(final BytecodeContext aBytecodeContext,
        final AnnotationValueBytecodeEvent aAnnotationValueBytecodeEvent) {
        if (inPaintsAnnotation && "value".equals(aAnnotationValueBytecodeEvent.getName())) {
            paintedType = (Type) aAnnotationValueBytecodeEvent.getValue();
        }
    }

    private void handleAnnotationEndBytecode(final BytecodeContext aBytecodeContext,
        final AnnotationEndBytecodeEvent aAnnotationEndBytecodeEvent) {
        inPaintsAnnotation = false;
    }

    private void handleMethodEndBytecode(final BytecodeContext aBytecodeContext,
        final MethodEndBytecodeEvent aMethodEndBytecodeEvent) {
        if (paintedType != null) {
            extendPaintComponent();
        }

        methodName = null;
    }

    private void handleClassEndBytecode(final BytecodeContext aBytecodeContext,
        final ClassEndBytecodeEvent aClassEndBytecodeEvent) {
        finishPaintComponent();
        aBytecodeContext.forwardBytecodeEvents(paintComponentBytecodeEvents);
    }

    private void preparePaintComponent() {
        paintComponentBytecodeEvents = new ArrayList<>();
        paintComponentBytecodeGenerator = Bytecode.createGenerator(paintComponentBytecodeEvents);

        paintComponentBytecodeGenerator.startMethod(Opcodes.ACC_PUBLIC, "paintComponent", PaintBrush.class,
            PaintContext.class);
        paintComponentBytecodeGenerator.startCode();

        paintComponentBytecodeGenerator.loadReference(1);
        paintComponentBytecodeGenerator.invokeInterface(PaintContext.class, "getComponent", Object.class);
        paintComponentBytecodeGenerator.storeReference(2);
    }

    private void extendPaintComponent() {
        final Label nextType = new Label();

        paintComponentBytecodeGenerator.pushConstant(paintedType);
        paintComponentBytecodeGenerator.loadReference(2);
        paintComponentBytecodeGenerator.invokeVirtual(Class.class, "isInstance", boolean.class, Object.class);
        paintComponentBytecodeGenerator.jumpEquals(nextType);

        paintComponentBytecodeGenerator.loadReference(0);
        paintComponentBytecodeGenerator.loadReference(1);
        paintComponentBytecodeGenerator.invokeVirtual(className, methodName, methodDescriptor);
        paintComponentBytecodeGenerator.returnReference();

        paintComponentBytecodeGenerator.markLabel(nextType);
    }

    private void finishPaintComponent() {
        paintComponentBytecodeGenerator.markFrameAppend(1, new Object[] { "java/lang/Object" }, 0, null);

        paintComponentBytecodeGenerator.loadReference(1);
        paintComponentBytecodeGenerator.invokeInterface(PaintContext.class, "forwardPaint", PaintBrush.class);
        paintComponentBytecodeGenerator.returnReference();

        paintComponentBytecodeGenerator.endCode(2, 3);
        paintComponentBytecodeGenerator.endMethod();
    }
}
